'''
Created on Feb 5, 2013

@author: nzhao
'''
import numpy as np
from traits.api import Array, List, Instance, Float

import phy.general.Constant as co
from phy.spin.MultiSpin import MultiSpin
from phy.interaction.SingleSpinInteraction import SingleSpinInteraction
from traits.has_traits import HasTraits

class tensor(HasTraits):
    xx, xy, xz = Float(), Float(), Float()
    yx, yy, yz = Float(), Float(), Float()
    zx, zy, zz = Float(), Float(), Float()
    
    def __init__(self, arr):
        self.xx, self.xy, self.xz = arr[0][0], arr[0][1], arr[0][2]  
        self.yx, self.yy, self.yz = arr[1][0], arr[1][1], arr[1][2]
        self.zx, self.zy, self.zz = arr[2][0], arr[2][1], arr[2][2]

class MultiSpinInteraction(MultiSpin):
    spin_list = List(SingleSpinInteraction)
    all_dipole_interaction = Array()
    
    def __init__(self, mode='list', spins=[], filename=None):
        MultiSpin.__init__(self, mode=mode, spins=spins, filename=filename)
        
    def dipolar(self, i, j, axes=None):
        spin1 = self.spin_list[i]
        spin2 = self.spin_list[j]
        if axes is None:
            axes = self.axes
        
        g1 = spin1.spin_relative_Gamma
        g2 = spin2.spin_relative_Gamma
        d = g1 * g2 * co.SpinSpinInteraction

        r_ij = self.dist(i, j)
        if r_ij == 0:
            res = [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]
        else:
            n_ij = self.dir_vect(i, j) / r_ij
            nx, ny, nz = np.dot(n_ij, axes.x_axis), np.dot(n_ij, axes.y_axis), np.dot(n_ij, axes.z_axis)
             
            res = d / r_ij**3 * np.array(
                    [[1.0 - 3.0 * nx*nx,      -3.0 * nx*ny,     -3.0 * nx*nz],\
                     [     -3.0 * ny*nx, 1.0 - 3.0 * ny*ny,     -3.0 * ny*nz],\
                     [     -3.0 * nz*nx,     - 3.0 * nz*ny, 1.0 -3.0 * nz*nz]])
        return res

    def dipolar_tensor(self, i, j, axes=None):
        arr = self.dipolar(i, j, axes)
        return tensor(arr)
    
    def dipolar_all(self, axes=None):
        self.all_dipole_interaction = [self.dipolar(i, j, axes) for i in range(self.nspin) for j in range(self.nspin) if i > j]
        
if __name__ == '__main__':
    s1=SingleSpinInteraction(element='C13', spin_type='Bath', x=0.0, y=0.0, z=0.0)
    s2=SingleSpinInteraction(element='C13', spin_type='Bath', x=1.0, y=1.0, z=1.0)
    s3=SingleSpinInteraction(element='C13', spin_type='Bath', x=1.0, y=3.0, z=5.0)
    
    spins = MultiSpinInteraction(spins=[s1,s2,s3])
    spins.dipolar_all()
    print spins.all_dipole_interaction
        